import json
import os
from typing import Dict, Any
from request_map_generator import MapGenerator3D, build_map_from_environmental_info
from multi_uav_planner_3D_waypoints import MultiUAVPlanner  
import pickle
from typing import List, Dict, Tuple, Optional
import csv


def load_map_data(filename):
    """加载保存的地图数据"""
    if os.path.exists(filename):  # 检查文件是否存在
        with open(filename, 'rb') as f:
            map_data = pickle.load(f)
        print(f"已加载地图数据: {filename}")
        print(f"地图尺寸: {map_data['x_width']}x{map_data['y_width']}x{map_data['z_width']}")
        return map_data
    else:
        return None  # 如果没有找到文件，返回 None


class MissionExecutor:
    def __init__(self, 
                 map_resolution: float = 20.0, 
                 robot_radius: float = 10.0, 
                 map_file: str = "test_map_from_env_info.pkl",
                 # 区域覆盖规划器参数（传递给MultiUAVPlanner）----------------
                 patrol_speed: float = 10.0,
                 min_patrol_speed: float = 5.0,
                 max_patrol_speed: float = 25.0,
                 target_patrol_time: float = 600.0,
                 time_buffer: float = 120.0,
                 default_altitude: float = 10.0,
                 sweep_width: float = 12.0,
                 max_segment_length: float = 20.0,
                 path_boundary_standoff: float = 5.0,
                 # 3D规划器参数（传递给MultiUAVPlanner）----------------
                 speed: float = 20.0,
                 safe_dis: float = 2.0,
                 min_planning_speed: float = 5.0,
                 max_planning_speed: float = 25.0,
                 speed_increase_factor: float = 1.1,
                 max_t: int = 3000,
                 goal_tolerance: float = 2.0,
                 heuristic_weight: float = 1.5):
        """
        任务执行器初始化，支持传递MultiUAVPlanner所需的所有关键参数
        
        Args:
            map_resolution: 地图分辨率（米/栅格）
            robot_radius: 无人机半径（用于避障计算）
            map_file: 预加载/保存的地图文件路径
            # 区域覆盖规划器参数（巡航、扫查任务用）
            patrol_speed: 巡航初始速度（m/s）
            min_patrol_speed: 巡航速度下限（m/s）
            max_patrol_speed: 巡航速度上限（m/s）
            target_patrol_time: 目标巡航时间（s）
            time_buffer: 任务时间缓冲（s）
            default_altitude: 默认飞行高度（m）
            sweep_width: 传感器扫描宽度（m）
            max_segment_length: 巡航路径段最大长度（m）
            path_boundary_standoff: 巡航路径与边界安全距离（m）
            # 3D规划器参数（路径搜索、避障用）
            speed: 3D规划初始速度（m/s）
            safe_dis: 与障碍物安全距离（m）
            min_planning_speed: 规划速度下限（m/s）
            max_planning_speed: 规划速度上限（m/s）
            speed_increase_factor: 速度调整系数（>1）
            max_t: 最大规划时间（s，超时判定失败）
            goal_tolerance: 目标点位置容差（m）
            heuristic_weight: A*启发式权重（1.1~2.0，值越大搜索越快）
        """
        # 基础参数
        self.map_resolution = map_resolution
        self.robot_radius = robot_radius
        self.map_file = map_file

        # 区域覆盖规划器参数
        self.patrol_speed = patrol_speed
        self.min_patrol_speed = min_patrol_speed
        self.max_patrol_speed = max_patrol_speed
        self.target_patrol_time=target_patrol_time
        self.time_buffer = time_buffer
        self.default_altitude = default_altitude
        self.sweep_width = sweep_width
        self.max_segment_length = max_segment_length
        self.path_boundary_standoff = path_boundary_standoff
        # 3D规划器参数
        self.speed = speed
        self.safe_dis = safe_dis
        self.min_planning_speed = min_planning_speed
        self.max_planning_speed = max_planning_speed
        self.speed_increase_factor = speed_increase_factor
        self.max_t = max_t
        self.goal_tolerance = goal_tolerance
        self.heuristic_weight = heuristic_weight
        
        # 加载地图并初始化规划器
        self.map_env = load_map_data(self.map_file)  # 尝试加载地图数据
        self.ref_lat = self.map_env['ref_lat'] if self.map_env else None
        self.ref_lon = self.map_env['ref_lon'] if self.map_env else None
        # 如果地图数据已加载，初始化规划器
        self.planner = self._init_planner() if self.map_env else None

    def _init_planner(self) -> MultiUAVPlanner:
        """内部工具方法：初始化MultiUAVPlanner"""
        if not self.map_env or self.ref_lat is None or self.ref_lon is None:
            raise ValueError("地图数据或参考经纬度缺失，无法初始化规划器")
        # 传递所有关键参数给MultiUAVPlanner
        return MultiUAVPlanner(
            map_env=self.map_env,
            ref_lat=self.ref_lat,
            ref_lon=self.ref_lon,
            # 区域覆盖规划器参数
            patrol_speed=self.patrol_speed,
            min_patrol_speed=self.min_patrol_speed,
            max_patrol_speed=self.max_patrol_speed,
            target_patrol_time=self.target_patrol_time,
            time_buffer=self.time_buffer,
            default_altitude=self.default_altitude,
            sweep_width=self.sweep_width,
            max_segment_length=self.max_segment_length,
            path_boundary_standoff=self.path_boundary_standoff,
            # 3D规划器参数
            speed=self.speed,
            safe_dis=self.safe_dis,
            min_planning_speed=self.min_planning_speed,
            max_planning_speed=self.max_planning_speed,
            speed_increase_factor=self.speed_increase_factor,
            max_t=self.max_t,
            goal_tolerance=self.goal_tolerance,
            heuristic_weight=self.heuristic_weight
        )

    # ---------------- 外部 JSON → 内部任务 ----------------
    def adapt_external_json(self, external_json: Dict[str, Any]):
        tasks = []

        # 1. point_mission → waypoint (支持途径点)
        if "point_mission" in external_json:
            pm = external_json["point_mission"]            
            uav_id = pm["deviceInfos"][0]["deviceId"]
            # 构建完整的航点序列：起点 + 途径点 + 终点
            waypoint_sequence = []
            
            # 起点
            waypoint_sequence.append({
                "lon": pm["startPoint"]["longitude"],
                "lat": pm["startPoint"]["latitude"], 
                "alt": pm["startPoint"]["altitude"]
            })
            
            # 途径点（按顺序）
            if "wayPoints" in pm and pm["wayPoints"]:
                for waypoint in pm["wayPoints"]:
                    waypoint_sequence.append({
                        "lon": waypoint["longitude"],
                        "lat": waypoint["latitude"],
                        "alt": waypoint["altitude"]
                    })
            
            # 终点
            waypoint_sequence.append({
                "lon": pm["endPoint"]["longitude"],
                "lat": pm["endPoint"]["latitude"],
                "alt": pm["endPoint"]["altitude"]
            })
            
            # 根据途径点数量决定任务类型，任务优先级暂定为1
            if len(waypoint_sequence) > 2:  # 有途径点
                task = {
                    "priority": 1,
                    "taskType": "multi_waypoint",
                    "uav_ids": [uav_id],
                    "time_constraint": pm["extraInfo"]["maxFlightTime"],
                    "start": [
                        {"uav_id": uav_id, "position": waypoint_sequence[0]}
                    ],
                    "waypoint_sequence": waypoint_sequence
                }
            else:  # 只有起点和终点
                task = {
                    "priority": 1,
                    "taskType": "waypoint",
                    "uav_ids": [uav_id],
                    "time_constraint": pm["extraInfo"]["maxFlightTime"],
                    "start": [
                        {"uav_id": uav_id, "position": waypoint_sequence[0]}
                    ],
                    "goals": [
                        {"uav_id": uav_id, "position": waypoint_sequence[1]}
                    ]
                }
            tasks.append(task)

        # 2. patrol_mission → patrol
        if "patrol_mission" in external_json:
            pm = external_json["patrol_mission"]
            uav_id = pm["deviceInfos"][0]["deviceId"]
            task = {
                "priority": 1,
                "taskType": "patrol",
                "uav_ids": [uav_id],
                "time_constraint": pm["extraInfo"]["maxFlightTime"],
                "start": [
                    {"uav_id": uav_id,
                     "position": {"lon": pm["startPoint"]["longitude"],
                                  "lat": pm["startPoint"]["latitude"],
                                  "alt": pm["startPoint"]["altitude"]}}
                ],
                "vertices": [
                    {"lon": v["longitude"], "lat": v["latitude"], "alt": v["altitude"]}
                    for v in pm["patrolArea"]["vertices"]
                ]
            }
            tasks.append(task)

        return tasks

    # ---------------- 整体执行流程 ----------------
    def execute(self, external_json: Dict[str, Any], map_file: str = "mission_map.pkl", show_visualization: bool = True):
        """
        执行任务规划
        
        Args:
            external_json: 外部任务JSON
            map_file: 地图文件名
            show_visualization: 是否显示可视化图表
            
        Returns:
            tuple: (gps_trajectories, world_paths) - GPS轨迹和世界坐标路径
        """
        # 1. 如果已有地图数据，跳过建图步骤
        if self.map_env:
            print("已存在地图数据，跳过建图。")
        else:
            # 2. JSON → 内部任务格式
            tasks = self.adapt_external_json(external_json)
            if not tasks:
                raise ValueError("外部 JSON 中没有可用任务")

            # 3. 只用第一个任务建图
            print(f"正在根据第一个任务建图: {tasks[0]['taskType']}")
            generated_map = build_map_from_environmental_info(
                beihang_request=external_json,
                map_resolution=self.map_resolution,
                robot_radius=self.robot_radius,
                output_filename=map_file
            )

            # 更新已有的地图数据
            self.map_env = load_map_data(map_file)
            self.ref_lat = self.map_env['ref_lat']
            self.ref_lon = self.map_env['ref_lon']
            # 重新初始化规划器（调用统一方法，传递参数）
            self.planner = self._init_planner()

        # 4. 把内部任务和地图传给多机规划器
        print("调用多机规划算法...")
        tasks = self.adapt_external_json(external_json)
        gps_trajectories, world_paths = self.planner.process_tasks(tasks)

        # ---------------- 给GPS轨迹添加建议速度 ----------------
        gps_with_speed = {}
        for uav_id in gps_trajectories.keys():
            if uav_id not in world_paths:
                print(f"无人机{uav_id}无世界路径，无法添加速度")
                gps_with_speed[uav_id] = gps_trajectories[uav_id]
                continue
            # 调用速度计算方法
            gps_with_speed[uav_id] = self.planner._add_speed_to_gps_trajectory(
                gps_traj=gps_trajectories[uav_id],
                world_path=world_paths[uav_id]
        )
        # 更新为带速度的GPS轨迹
        gps_trajectories = gps_with_speed
        # 可视化
        if show_visualization and world_paths:
            print("\n开始可视化...")
            try:
                self.planner.plot_static_world_trajectories_3d(world_paths)
                self.planner.animate_world_trajectories_3d(world_paths, interval=100)
            except Exception as e:
                print(f"可视化过程中出现错误: {e}")

        return gps_trajectories, world_paths


if __name__ == "__main__":
    # 1. 加载外部任务JSON
    with open("test_mission_data.json", "r", encoding="utf-8") as f:
        external_json = json.load(f)

    # 2. 初始化任务执行器
    executor = MissionExecutor(
        map_resolution=15.0,          # 地图分辨率提高
        robot_radius=8.0,             # 无人机半径减小
        map_file="test_map_from_env_info.pkl",
        # ---------------- 自定义区域覆盖规划器参数 ----------------
        patrol_speed=15.0,            # 巡航速度提高到15m/s
        target_patrol_time=900.0,     # 目标巡航时间延长到15分钟
        time_buffer=180.0,            # 时间缓冲延长到3分钟
        default_altitude=15.0,        # 默认飞行高度提高到15m
        sweep_width=40.0,             # 传感器扫描宽度扩大到40m
        path_boundary_standoff=10.0,   # 与边界安全距离扩大到10m
        # ---------------- 自定义3D规划器参数 ----------------
        speed=20.0,                   # 3D规划初始速度提高到20m/s
        safe_dis=10.0,                 # 与障碍物安全距离扩大到10m
        heuristic_weight=1.8,         # 启发式权重提高（优先搜索速度）
        max_t=3600                    # 最大规划时间延长到1小时
    )
    
    # 3. 执行规划并自动显示可视化
    gps_result, world_result = executor.execute(
        external_json, 
        map_file="test_map_from_env_info.pkl", 
        show_visualization=True  # 设置为False可禁用可视化
    )

    # 4. 输出规划结果到CSV
    print("\n规划结果:")
    if gps_result:
        # 创建输出目录（如果不存在）
        output_dir = "mission_results"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 为每个无人机生成CSV文件
        for uav_id, trajectory in gps_result.items():
            filename = os.path.join(output_dir, f"UAV_{uav_id}_trajectory.csv")
            
            with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                # 写入表头
                writer.writerow(['Latitude', 'Longitude', 'Altitude', 'Speed (m/s)'])
                # 写入航点数据
                for waypoint in trajectory:
                    writer.writerow([
                        waypoint['lat'],
                        waypoint['lon'],
                        waypoint['alt'],
                        waypoint.get('speed', 0.0)  # 若速度缺失，默认0
                    ])
            
            print(f"无人机 {uav_id} 的航点已保存到: {filename}")
            print(f"  共 {len(trajectory)} 个航点")